﻿using CocoaAni.Net.WebSockets.Enums;
using CocoaAni.Net.WebSockets.Exceptions;
using CocoaAni.Net.WebSockets.Frames;
using CocoaAni.Net.WebSockets.Messages;
using CocoaAni.Net.WebSockets.Tools;
using System.Net;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;
using System.Text.RegularExpressions;

namespace CocoaAni.Net.WebSockets;

public delegate Task HandShakeHandler(WSHandShakeRequest request, WSHandShakeResponse response);

public class WebSocket
{
    public WebSocket(Socket socket)
    {
        Socket = socket;
    }

    public Uri? Url { get; protected set; }
    public Socket? Socket { get; protected set; }
    public bool IsHandShake { get; protected set; }
    public Stream BaseStream { get; set; } = null!;
    public bool IsSecurity { get; set; }
    public byte[]? Masks { get; set; }
    public EndPoint? RemoteEndPoint => Socket?.RemoteEndPoint;
    public EndPoint? LocalEndPoint => Socket?.LocalEndPoint;

    public async Task InitForClientAsync(Uri url,
        Action<WSHandShakeRequest>? requestInitializer = null,
        Action<WSHandShakeResponse>? onResponse = null,
        CancellationToken cancellationToken = default)
    {
        Url = url;
        var addresses = await Dns.GetHostAddressesAsync(Url!.Host, cancellationToken);
        if (addresses.Length == 0)
            throw new WebSocketException($"Host IpAddress Is Not Found！=>{Url.Host}");
        Socket = new Socket(addresses[0].AddressFamily, SocketType.Stream, ProtocolType.IP);
        await Socket.ConnectAsync(addresses[0], Url.Port, cancellationToken);
        if (!Socket.Connected)
        {
            throw new WebSocketException($"Socket Connect To {addresses[0]} Failed!");
        }
        BaseStream = new NetworkStream(Socket);
        try
        {
            if (IsSecurity)
            {
                var sslStream = typeof(SslStream) == BaseStream.GetType() ? (SslStream)BaseStream : new SslStream(BaseStream);
                await sslStream.AuthenticateAsClientAsync(Url.Host);
                BaseStream = sslStream;
            }

            var request = new WSHandShakeRequest();
            requestInitializer?.Invoke(request);
            //var handshakeData = CreateHandShakeRequest(HandShakeKey);
            await BaseStream.WriteAsync(request.Build(), cancellationToken);
            var respDataBuf = new byte[4096];
            var respDataLen = await BaseStream.ReadAsync(respDataBuf.AsMemory(), cancellationToken);
            var resp = new WSHandShakeResponse().Parse(respDataBuf[..respDataLen]);
            onResponse?.Invoke(resp);
            IsHandShake = resp.StatusCode != 101;
        }
        catch (Exception e)
        {
            throw new WebSocketException($"WebSocket Init For Client Error: {e.Message}", e); ;
        }
    }

    public async Task InitForServerSideAsync(
        HandShakeHandler? handShakeHandler,
        Func<Stream, Task<SslStream>>? securityStreamInitializer = null,
        CancellationToken cancellationToken = default)
    {
        BaseStream = new NetworkStream(Socket!);
        if (securityStreamInitializer != null)
        {
            BaseStream = await securityStreamInitializer.Invoke(BaseStream);
        }

        var buffer = new byte[4096];
        var readLength = await BaseStream.ReadAsync(buffer, 0, buffer.Length, cancellationToken);
        if (handShakeHandler == null)
        {
            await BaseStream.WriteAsync(CreateHandShakeResponse(buffer, readLength), cancellationToken);
            IsHandShake = true;
            return;
        }
        var request = new WSHandShakeRequest().Parse(buffer, readLength);
        var response = request.CreateOkResponse();
        await handShakeHandler.Invoke(request, response);
        await BaseStream.WriteAsync(response.Build(), cancellationToken);
        IsHandShake = response.StatusCode == 101;
        if (!IsHandShake)
        {
            throw new WebSocketException("Handle Shake Filed!");
        }
    }

    private byte[] CreateHandShakeResponse(byte[] request, int length)
    {
        var handShakeText = Encoding.UTF8.GetString(request, 0, length);

        var key = string.Empty;
        var matched = new Regex(@"Sec\-WebSocket\-Key:(.*?)\r\n").Match(handShakeText);
        if (matched.Groups.Count != 0)
        {
            key = Regex.Replace(matched.Value, @"Sec\-WebSocket\-Key:(.*?)\r\n", "$1").Trim();
        }
        var encryptionString = SHA1.Create().ComputeHash(Encoding.ASCII.GetBytes(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"));
        var respText =
            $"HTTP/1.1 101 Switching Protocols{Environment.NewLine}" +
            $"Upgrade: websocket{Environment.NewLine}" +
            $"Connection: Upgrade{Environment.NewLine}" +
            $"Sec-WebSocket-Accept: {Convert.ToBase64String(encryptionString)}{Environment.NewLine}{Environment.NewLine}";
        return Encoding.UTF8.GetBytes(respText);
    }

    protected byte[] CreateHandShakeRequest(string key, WSHandShakeRequest? request = null)
    {
        var origin = $"{(Url!.Scheme == "wss" ? "https://" : "http://")}{Url.Host}:{Url.Port}";
        if (request != null)
        {
            if (request.Headers.ContainsKey("Sec-WebSocket-Key"))
            {
                request.Headers["Sec-WebSocket-Key"] = key;
            }
            request.Headers["Host"] = Socket!.LocalEndPoint!;
            request.Headers["Origin"] = origin;
            return request.Build();
        }
        var requestText =
            $"GET / HTTP/1.1{Environment.NewLine}" +
            $"Host: {Socket!.LocalEndPoint}{Environment.NewLine}" +
            $"Origin: {origin}{Environment.NewLine}" +
            $"Connection: Upgrade{Environment.NewLine}" +
            $"Upgrade: websocket{Environment.NewLine}" +
            $"Sec-WebSocket-Version: 13{Environment.NewLine}" +
            $"Sec-WebSocket-Key: {key}{Environment.NewLine}{Environment.NewLine}";
        return Encoding.UTF8.GetBytes(requestText);
    }

    public Message Receive(int bufferSize) => ReceiveAsync(bufferSize).Result;

    public Message Receive(ReceiveBuffer buffer) => ReceiveAsync(buffer).Result;

    public string ReceiveText(int bufferSize = 40960) => ReceiveTextAsync(bufferSize).Result;

    public Task<Message> ReceiveAsync(int bufferSize, CancellationToken cancellationToken = default)
    {
        var buffer = new ReceiveBuffer(new byte[bufferSize], 0, false);
        return ReceiveAsync(buffer, cancellationToken);
    }

    public async Task<Message> ReceiveAsync(ReceiveBuffer buffer, CancellationToken cancellationToken = default)
    {
        while (true)
        {
            var frame = await ReceiveFrameAsync(buffer, cancellationToken);
            if (frame == null) continue;
            if (frame.IsFinal)
            {
                return buffer.BuildMessage();
            }
        }
    }

    public async Task<string> ReceiveTextAsync(int bufferSize = 40960, CancellationToken cancellationToken = default)
    {
        var buffer = new ReceiveBuffer(bufferSize)
        {
            IsCopyDataOnBuildMessage = false,
            IsSaveHeader = false,
        };
        var message = await ReceiveAsync(buffer, cancellationToken);
        return Encoding.UTF8.GetString(message.DataMemory.Span);
    }

    public void Send(byte[] buffer, int offset, int count) => SendAsync(buffer, offset, count).Wait();

    public void SendText(string text) => SendTextAsync(text).Wait();

    public void SendText(string text, int split) => SendTextAsync(text, split).Wait();

    public Task SendAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default)
        => BaseStream.WriteAsync(buffer, offset, count, cancellationToken);

    public Task SendTextAsync(string text, CancellationToken cancellationToken = default)
        => SendFrameAsync(Encoding.UTF8.GetBytes(text), SendFrameBuildMode.Split, true, WSOpcode.Text, cancellationToken).AsTask();

    public async Task SendTextAsync(string text, int splitSize, CancellationToken cancellationToken = default)
    {
        var bytes = Encoding.UTF8.GetBytes(text);
        var sendCounter = 0;
        while (true)
        {
            var sendSize = sendCounter + splitSize > bytes.Length ? bytes.Length - sendCounter : splitSize;
            await SendFrameAsync(bytes, sendCounter, sendSize,
                SendFrameBuildMode.Split,
                sendSize != splitSize,
                sendCounter == 0 ? WSOpcode.Text : WSOpcode.Continue,
                cancellationToken);
            sendCounter += sendSize;
            if (sendSize != splitSize)
                return;
        }
    }

    public Task SendTextFileAsync(string path, CancellationToken cancellationToken = default)
        => SendFileAsync(path, 0, true, null, cancellationToken);

    public Task SendTextFileAsync(string path, int split, CancellationToken cancellationToken = default)
        => SendFileAsync(path, split, true, null, cancellationToken);

    public Task SendFileAsync(string path, CancellationToken cancellationToken = default)
        => SendFileAsync(path, 0, false, null, cancellationToken);

    public Task SendFileAsync(string path, int split, CancellationToken cancellationToken = default)
        => SendFileAsync(path, split, false, null, cancellationToken);

    public struct SendFileUpdateInfo
    {
        public string Path;
        public long FileSize;
        public long SendSize;
        public int SplitSize;
    }

    public async Task SendFileAsync(string path, int split, bool isText, Action<SendFileUpdateInfo>? onUpdate, CancellationToken cancellationToken = default)
    {
        var sendInfo = new SendFileUpdateInfo();
        var fileStream = File.OpenRead(path);
        sendInfo.SplitSize = split;
        sendInfo.Path = path;
        sendInfo.FileSize = fileStream.Length;
        if (split <= 0)
        {
            if (sendInfo.FileSize > int.MaxValue)
            {
                throw new WebSocketException("FileSize > int.MaxValue,You Can Set The split");
            }
            split = (int)sendInfo.FileSize;
        }
        var hasMask = Masks != null;
        var headerLength = Frame.CalculateHeaderSize(hasMask, split);
        var buffer = new byte[headerLength + split];

        var frame = new SendFrame()
        {
            Buffer = buffer,
            HeaderBuffer = buffer,
            HasMask = hasMask,
            PayloadLength = split,
            Masks = Masks,
            Opcode = isText ? WSOpcode.Text : WSOpcode.Binary,
            HeaderStartIndex = 0,
            PayloadStartIndex = headerLength,
            //PayloadEndIndex = buffer.Length
        };
        while (true)
        {
            var readLength = await fileStream.ReadAsync(buffer.AsMemory(headerLength), cancellationToken);
            if (readLength is 0)
                break;
            if (sendInfo.SendSize + readLength == sendInfo.FileSize)
            {
                frame.IsFinal = true;
            }
            frame.UpdatePayloadLength(readLength);
            await SendFrameAsync(frame, cancellationToken);
            sendInfo.SendSize += readLength;
            frame.Opcode = WSOpcode.Continue;
            onUpdate?.Invoke(sendInfo);
        }
        fileStream.Close();
    }

    public void Close()
    {
        BaseStream.Close();
        Socket?.Close();
    }

    public async Task<Frame?> ReceiveFrameAsync(ReceiveBuffer buffer, CancellationToken cancellationToken = default)
    {
        if (!IsHandShake)
        {
            throw new WebSocketException("Please Call InitAsync To HandShake!");
        }
        while (true)
        {
            var handleSize = 0;
            var readLength = await BaseStream.ReadAsync(buffer.Buffer, buffer.HandleCursor, buffer.BufferCanWriteCount, cancellationToken);
            if (readLength == 0) continue;
            if (readLength == -1) break;
            while (handleSize != readLength)
            {
                if (buffer.PayloadLength > buffer.BufferCanWriteCount)
                {
                    throw new WebSocketException
                        ($"PayloadLength={buffer.PayloadLength} > BufferCanWriteCount={buffer.BufferCanWriteCount}");
                }
                handleSize += buffer.AppendData(readLength);
                if (buffer.IsComplete)
                {
                    return buffer.BuildFrame();
                }
            }
        }
        return buffer.BuildFrame();
    }

    public Task<Frame?> ReceiveFrameAsync(int bufSize, CancellationToken cancellationToken = default)
    {
        var buffer = new ReceiveBuffer(new byte[bufSize], 0, true);
        return ReceiveFrameAsync(buffer, cancellationToken);
    }

    public async ValueTask SendFrameAsync(SendFrame frame, CancellationToken cancellationToken = default)
    {
        if (!IsHandShake)
        {
            throw new WebSocketException("Please Call InitAsync To HandShake!");
        }
        if (frame.Buffer == frame.HeaderBuffer)
        {
            await BaseStream.WriteAsync(frame.Buffer.AsMemory(frame.HeaderStartIndex, frame.HeaderLength + frame.PayloadLength), cancellationToken);
            return;
        }
        await BaseStream.WriteAsync(frame.HeaderBuffer, cancellationToken);
        await BaseStream.WriteAsync(frame.Payload, cancellationToken);
    }

    public ValueTask SendFrameAsync(byte[] buffer, int offset, int length, SendFrameBuildMode saveMode, bool isFin, WSOpcode opcode, CancellationToken cancellationToken = default)
    {
        var frame = new SendFrame(buffer, offset, length, saveMode, isFin, opcode, Masks);
        return SendFrameAsync(frame, cancellationToken);
    }

    public ValueTask SendFrameAsync(byte[] buffer, SendFrameBuildMode saveMode, bool isFin, WSOpcode opcode, CancellationToken cancellationToken = default)
    {
        var frame = new SendFrame(buffer, 0, buffer.Length, saveMode, isFin, opcode, null);
        return SendFrameAsync(frame, cancellationToken);
    }
}